"""
数据加载和处理模块
"""

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Dataset
from config import Config

def get_cifar100_superclass_mapping():
    """
    返回CIFAR100超类到类的映射关系
    每个超类包含多个细粒度类
    """
    return {
        0: ['beaver', 'dolphin', 'otter', 'seal', 'whale'],
        1: ['aquarium_fish', 'flatfish', 'ray', 'shark', 'trout'],
        2: ['orchid', 'poppy', 'rose', 'sunflower', 'tulip'],
        3: ['bottle', 'bowl', 'can', 'cup', 'plate'],
        4: ['apple', 'mushroom', 'orange', 'pear', 'sweet_pepper'],
        5: ['clock', 'keyboard', 'lamp', 'telephone', 'television'],
        6: ['bed', 'chair', 'couch', 'table', 'wardrobe'],
        7: ['bee', 'beetle', 'butterfly', 'caterpillar', 'cockroach'],
        8: ['bear', 'leopard', 'lion', 'tiger', 'wolf'],
        9: ['bridge', 'castle', 'house', 'road', 'skyscraper'],
        10: ['cloud', 'forest', 'mountain', 'plain', 'sea'],
        11: ['camel', 'cattle', 'chimpanzee', 'elephant', 'kangaroo'],
        12: ['fox', 'porcupine', 'possum', 'raccoon', 'skunk'],
        13: ['crab', 'lobster', 'snail', 'spider', 'worm'],
        14: ['baby', 'boy', 'girl', 'man', 'woman'],
        15: ['crocodile', 'dinosaur', 'lizard', 'snake', 'turtle'],
        16: ['hamster', 'mouse', 'rabbit', 'shrew', 'squirrel'],
        17: ['maple_tree', 'oak_tree', 'palm_tree', 'pine_tree', 'willow_tree'],
        18: ['bicycle', 'bus', 'motorcycle', 'pickup_truck', 'train'],
        19: ['lawn_mower', 'rocket', 'streetcar', 'tank', 'tractor']
    }

def get_transforms():
    """
    返回训练和测试的数据增强/预处理
    """
    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
    ])
    
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
    ])
    
    return train_transform, test_transform

def get_dataloaders():
    """
    获取CIFAR100数据加载器(超类级别)
    """
    cfg = Config()
    
    # 获取数据变换
    train_transform, test_transform = get_transforms()
    
    # 加载数据集
    train_dataset = datasets.CIFAR100(
        root='./data', 
        train=True, 
        download=True, 
        transform=train_transform
    )
    
    test_dataset = datasets.CIFAR100(
        root='./data', 
        train=False, 
        download=True, 
        transform=test_transform
    )
    
    # 将细粒度标签转换为超类标签
    superclass_mapping = {}
    for superclass_idx, classes in get_cifar100_superclass_mapping().items():
        for class_name in classes:
            class_idx = train_dataset.class_to_idx[class_name]
            superclass_mapping[class_idx] = superclass_idx
    
    # 转换标签函数
    def convert_to_superclass(dataset):
        data = []
        targets = []
        for img, target in dataset:
            data.append(img)
            targets.append(superclass_mapping[target])
        return list(zip(data, targets))
    
    # 转换数据集标签
    train_dataset = convert_to_superclass(train_dataset)
    test_dataset = convert_to_superclass(test_dataset)
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset, 
        batch_size=cfg.BATCH_SIZE, 
        shuffle=True, 
        num_workers=cfg.NUM_WORKERS,
        pin_memory=True
    )
    
    test_loader = DataLoader(
        test_dataset, 
        batch_size=cfg.BATCH_SIZE, 
        shuffle=False, 
        num_workers=cfg.NUM_WORKERS,
        pin_memory=True
    )
    
    return train_loader, test_loader